#pragma once

#include "tree.hpp"

enum class Color { RED, BLACK, DOUBLE_BLACK };
enum class PosType { LL, LR, RL, RR, Others };

template <typename K, typename V>
class RBT : public Tree<K, V> {
  using Tree<K, V>::is_nil;
  using Tree<K, V>::type;

 private:
  RBT<K, V>* left_rotate() override {
    return dynamic_cast<RBT<K, V>*>(Tree<K, V>::left_rotate());
  }
  RBT<K, V>* right_rotate() override {
    return dynamic_cast<RBT<K, V>*>(Tree<K, V>::right_rotate());
  }

  // 父节点
  RBT<K, V>* get_p() const { return dynamic_cast<RBT<K, V>*>(this->p); }
  // 左孩子
  RBT<K, V>* get_l() const { return dynamic_cast<RBT<K, V>*>(this->l); }
  // 右孩子
  RBT<K, V>* get_r() const { return dynamic_cast<RBT<K, V>*>(this->r); }
  // 兄弟节点
  RBT<K, V>* get_b() const {
    if (type == NodeType::ROOT) {
      return nullptr;
    } else if (type == NodeType::LEFT) {
      return get_p()->get_r();
    } else {
      return get_p()->get_l();
    }
  }
  // 近侄子
  RBT<K, V>* get_close_nephew() const {
    if (type == NodeType::ROOT) {
      return nullptr;
    } else if (type == NodeType::LEFT) {
      return get_b()->get_l();
    } else {
      return get_b()->get_r();
    }
  }
  // 远侄子
  RBT<K, V>* get_distant_nephew() const {
    return get_close_nephew() ? get_close_nephew()->get_b() : nullptr;
  }

  Color color{Color::BLACK};
  mutable int bh{0};
  mutable bool bh_valid{true};

  // 每个节点或是红色，或是黑色
  bool principle_1() const {
    if (is_nil) {
      return color != Color::DOUBLE_BLACK;
    } else {
      return get_l()->principle_1() && get_r()->principle_1() &&
             color != Color::DOUBLE_BLACK;
    }
  }
  // 根节点是黑色的
  bool principle_2() const {
    const RBT<K, V>* root = this;
    while (root->get_p()) {
      root = root->get_p();
    }
    return root->color == Color::BLACK;
  }
  // 叶节点是黑色的
  bool principle_3() const {
    if (is_nil) {
      return color == Color::BLACK;
    } else {
      return get_l()->principle_3() && get_r()->principle_3();
    }
  }
  // 不存在两个相邻的红节点
  bool principle_4() const {
    if (is_nil) {
      return true;
    } else {
      return (color != Color::RED ||
              (get_l()->color != Color::RED && get_r()->color != Color::RED)) &&
             get_l()->principle_4() && get_r()->principle_4();
    }
  }
  // 黑高均等
  bool principle_5() const {
    if (is_nil) {
      bh_valid = true;
      bh = 0;
    } else {
      bh_valid = get_l()->principle_5() && get_r()->principle_5() &&
                 ((get_l()->black_cnt() + get_l()->bh) ==
                  (get_r()->black_cnt() + get_r()->bh));
      bh = get_l()->black_cnt() + get_l()->bh;
    }
    return bh_valid;
  }
  // 黑值计算
  int black_cnt() const {
    if (color == Color::RED) {
      return 0;
    } else if (color == Color::BLACK) {
      return 1;
    } else {
      return 2;
    }
  }

  // 获取this的两层位置类型
  PosType get_pos_type() const {
    if (!get_p() || !(get_p()->get_p())) {
      return PosType::Others;
    } else {
      if (get_p()->type == NodeType::LEFT) {
        if (type == NodeType::LEFT) {
          return PosType::LL;
        } else {
          return PosType::LR;
        }
      } else {
        if (type == NodeType::LEFT) {
          return PosType::RL;
        } else {
          return PosType::RR;
        }
      }
    }
  }

  // 为节点增添一层黑色
  void add_black() {
    assert(color != Color::DOUBLE_BLACK);
    if (color == Color::RED) {
      color = Color::BLACK;
    } else {
      color = Color::DOUBLE_BLACK;
    }
  }

  // 消除双黑属性的函数
  // 由可能为双黑的节点调用
  // 注: 初代双黑必为nil，但经过向上传递可能出现非nil双黑
  void eliminate_db() {
    // 非双黑节点不管
    if (color != Color::DOUBLE_BLACK) return;
    RBT<K, V>* to_adjust = this;
    while (to_adjust) {
      int cond = to_adjust->classify_db();
      to_adjust = to_adjust->handle_db(cond);
    }
  }
  // 双黑情况分类机
  int classify_db() const {
    assert(color == Color::DOUBLE_BLACK);
    if (type == NodeType::ROOT) {
      // 为根
      return 0;
    } else if (get_b()->color == Color::RED) {
      // 兄弟红
      return 1;
    } else if (get_distant_nephew()->color == Color::RED) {
      // 兄弟黑，远侄子红
      return 2;
    } else if (get_close_nephew()->color == Color::RED) {
      // 兄弟黑，远侄子黑，近侄子红
      return 3;
    } else {
      //// 兄弟黑，侄子均黑
      return 4;
    }
  }
  // 双黑情况处理机
  // 返回新双黑节点或者残余指针
  RBT<K, V>* handle_db(int condition) {
    assert(color == Color::DOUBLE_BLACK);
    RBT<K, V>* new_or_remain_db = nullptr;
    switch (condition) {
      case 0: {
        color = Color::BLACK;
        break;
      }
      case 1: {
        RBT<K, V>* temp = nullptr;
        if (type == NodeType::LEFT) {
          temp = get_p()->left_rotate();
          temp->color = Color::BLACK;
          temp->get_l()->color = Color::RED;
        } else {
          temp = get_p()->right_rotate();
          temp->color = Color::BLACK;
          temp->get_r()->color = Color::RED;
        }
        new_or_remain_db = this;
        break;
      }
      case 2: {
        RBT<K, V>* temp = nullptr;
        if (type == NodeType::LEFT) {
          temp = get_p()->left_rotate();
          color = Color::BLACK;
          temp->color = get_p()->color;
          get_p()->color = Color::BLACK;
          temp->get_r()->color = Color::BLACK;
        } else {
          temp = get_p()->right_rotate();
          color = Color::BLACK;
          temp->color = get_p()->color;
          get_p()->color = Color::BLACK;
          temp->get_l()->color = Color::BLACK;
        }
        break;
      }
      case 3: {
        if (type == NodeType::LEFT) {
          get_b()->right_rotate();
          get_b()->color = Color::BLACK;
          get_distant_nephew()->color = Color::RED;
        } else {
          get_b()->left_rotate();
          get_b()->color = Color::BLACK;
          get_distant_nephew()->color = Color::RED;
        }
        new_or_remain_db = this;
        break;
      }
      case 4: {
        color = Color::BLACK;
        get_b()->color = Color::RED;
        get_p()->add_black();
        new_or_remain_db =
            get_p()->color == Color::DOUBLE_BLACK ? get_p() : nullptr;
        break;
      }
      default: {
        break;
      }
    }
    return new_or_remain_db;
  }

 protected:
  Tree<K, V>* new_nil() override { return new RBT(); }
  void post_insert() override {
    // 插入节点先默认红色
    color = Color::RED;
    if (!get_p()) {
      // 如果是根节点，变为黑色
      color = Color::BLACK;
      return;
    } else if (get_p()->color == Color::BLACK) {
      // 如果父节点黑，则未影响红黑树性质，无需调整
      return;
    } else {
      // 此时父节点红，则必有爷节点叔节点
      RBT<K, V>* new_red = this;
      RBT<K, V>* parent = get_p();
      RBT<K, V>* grandpa = parent->get_p();
      RBT<K, V>* uncle = parent->get_b();
      if (uncle->color == Color::BLACK) {
        // 此时新红节点位置为LL/LR/RL/RR之一
        // LR/RL先转化为LL/RR
        if (new_red->get_pos_type() == PosType::LR) {
          new_red = parent;
          parent = parent->left_rotate();
        }
        if (new_red->get_pos_type() == PosType::RL) {
          new_red = parent;
          parent = parent->right_rotate();
        }
        if (new_red->get_pos_type() == PosType::LL) {
          grandpa = grandpa->right_rotate();
          grandpa->color = Color::BLACK;
          grandpa->get_l()->color = Color::RED;
          grandpa->get_r()->color = Color::RED;
          return;
        }
        if (new_red->get_pos_type() == PosType::RR) {
          grandpa = grandpa->left_rotate();
          grandpa->color = Color::BLACK;
          grandpa->get_l()->color = Color::RED;
          grandpa->get_r()->color = Color::RED;
          return;
        }

      } else {
        // 叔节点为红
        grandpa->color = Color::RED;
        parent->color = Color::BLACK;
        uncle->color = Color::BLACK;
        grandpa->post_insert();
      }
    }
  }
  void post_remove() override {
    if (!is_nil) {
      // 继任节点非nil说明原最终删除节点有左或者由子树
      // 则原子树根必红，则原最终删除节点必黑
      // 继任节点保持原最终删除节点颜色，无需调整
      return;
    } else {
      // 继任节点为nil说明原最终删除节点无孩子，从普通节点退为nil
      // 需补偿一层黑色以抵消移除原nil的黑高变化
      add_black();
      eliminate_db();
    }
  }
  bool extra_check() const override {
    bool ok_1 = principle_1();
    if (!ok_1) {
      std::cout << "不满足红黑性质1" << std::endl;
    }
    bool ok_2 = principle_2();
    if (!ok_2) {
      std::cout << "不满足红黑性质2" << std::endl;
    }
    bool ok_3 = principle_3();
    if (!ok_3) {
      std::cout << "不满足红黑性质3" << std::endl;
    }
    bool ok_4 = principle_4();
    if (!ok_4) {
      std::cout << "不满足红黑性质4" << std::endl;
    }
    bool ok_5 = principle_5();
    if (!ok_5) {
      std::cout << "不满足红黑性质5" << std::endl;
    }
    return ok_1 && ok_2 && ok_3 && ok_4 && ok_5;
  }
  const char* extra_info() const override {
    if (color == Color::RED) {
      return "(R)";
    } else if (color == Color::BLACK) {
      return "(B)";
    } else {
      return "(D)";
    }
  }
};
